from .revolution import qaoa, dc_qaoa, abf_qaoa
from mindquantum.core.circuit import Circuit, UN
from mindquantum.core.gates import H, Rzz, RX, RZ, RY
from mindquantum.core.operators import Hamiltonian, QubitOperator
from mindquantum.framework import MQAnsatzOnlyLayer
from mindquantum.simulator import Simulator
import networkx as nx
import mindspore.nn as nn
from scipy.sparse import csr_matrix
import scipy.sparse.linalg as sla
import numpy as np
import mindspore as ms
from utils.revolution import qaoa, dc_qaoa, abf_qaoa
from utils.ham import get_ham_matrix

ms.set_context(mode=ms.PYNATIVE_MODE, device_target="CPU")

class QAOA:
    def __init__(self, config: dict):
        self.config = config
        self._set_circuit_type()

    def _set_circuit_type(self):
        self.circuit_type = self.config['circuit_type']
        self.N = self.config['N']
        self.depth = self.config['depth']
        self.P = self.config['P']
        self.ideal_state = self.config['ideal_state']
        self.add_noise = self.config['add_noise']
        self.noise_param = self.config['noise_param']

        if self.circuit_type == 'qaoa':
            self.circ = qaoa(self.N, self.depth)
        elif self.circuit_type == 'dc_qaoa':
            self.circ = dc_qaoa(self.N, self.depth)
        elif self.circuit_type == 'abf_qaoa':
            self.circ = abf_qaoa(self.N, self.depth)
        if self.add_noise:
            from mindquantum.core.circuit import MixerAdder, DepolarizingChannelAdder, QubitNumberConstrain
            # we add depolarizing noise to the circuit
            p1, p2 = self.noise_param
            noise_adder1 = MixerAdder([QubitNumberConstrain(1),DepolarizingChannelAdder(p1, 1),], add_after=True)
            noise_adder2 = MixerAdder([QubitNumberConstrain(2),DepolarizingChannelAdder(p2, 2),], add_after=True)

            self.circ = noise_adder2(noise_adder1(self.circ))

    def _build_ham(self) -> np.ndarray:
        ham_matrix = get_ham_matrix(self.config['BIG_adjacent_weight_matrix'], self.config['BIG_adjacent_matrix'], self.config['P'])
        return ham_matrix

    def evolution(self, print_ratio: bool = False) -> np.ndarray:
        lr = self.config['lr']
        epoch = self.config['epoch']
        circ = self.circ

        # if self.circuit_type == 'qaoa' or self.circuit_type == 'dc_qaoa':    
        sim = Simulator('mqvector', circ.n_qubits)                     # 创建模拟器，backend使用‘mqvector’，能模拟5个比特（'circ'线路中包含的比特数）
        
        ham_matrix = self._build_ham()
        ham = csr_matrix(ham_matrix)

        min_eigenvalue, min_eigenvector = sla.eigsh(ham, k=1, which='SA')
        max_eigenvalue, max_eigenvector = sla.eigsh(ham, k=1, which='LA')

        # successful_ratios = []

        grad_ops = sim.get_expectation_with_grad(Hamiltonian(ham), circ)            # 获取计算变分量子线路的期望值和梯度的算子
        net = MQAnsatzOnlyLayer(grad_ops)                              # 生成待训练的神经网络
        opti = nn.Adam(net.trainable_params(), learning_rate=lr)     # 设置针对网络中所有可训练参数、学习率为lr的Adam优化器
        train_net = nn.TrainOneStepCell(net, opti)

        self.training_data = []
        from tqdm import tqdm  # Add tqdm for progress bar

        for e in tqdm(range(epoch), desc="Training Progress"):
            train_net()
            thetas = net.weight.asnumpy()
            f, _ = grad_ops(thetas)
            ratio = (max_eigenvalue-f.real)/(max_eigenvalue-min_eigenvalue)
            self.training_data.append(ratio[0][0])
            if print_ratio and e % 10 == 0:
                print(f'epoch: {e}, approximate ratio: {ratio[0][0]}')
        
        print(f'final approximate ratio: {ratio[0][0]}')
        print(self.circuit_type + ' parameter optimization finished')
        
        self.thetas = thetas
        # return thetas
    
    def get_successful_ratio(self, shots: int, show_result: bool = False) -> float:
        successful_ratio = 0.0
        
        circ = self.circ
        circ.measure_all()
        sim = Simulator('mqvector', circ.n_qubits)
        
        pr = dict(zip(circ.params_name, self.thetas))
        result = sim.sampling(circ, pr=pr, shots=shots).data # 比特串采用小端序表示(返回字典)
        if show_result:
            print(result)
        if self.ideal_state in result:
            successful_ratio = result[self.ideal_state]/shots
        else:
            print(f'ideal state {self.ideal_state} not found in result')
        
        return successful_ratio

    def get_training_data(self) -> np.ndarray:
        return np.array(self.training_data)

class ABF_QAOA:
    def __init__(self, config):
        self.config = config
        self._set_circuit()
    
    def _set_circuit(self):
        self.N = self.config['N']
        self.depth = self.config['depth']
        self.P = self.config['P']
        self.ideal_state = self.config['ideal_state']
        self.add_noise = self.config['add_noise']
        self.noise_param = self.config['noise_param']

        self.init_part = Circuit()
        # self.init_part += UN(H, self.N)
        for i in range(self.N):
            # self.init_part += RZ(f'init_bias_{i}').on(i)
            self.init_part += RY(f'init_bias_{i}').on(i)
    

        self.train_part = Circuit()
        for p in range(self.depth):
            # generate U_C
            for i in range(self.N-3):
                self.train_part += Rzz(f'rzz_{p}{i}').on((i, i+3))
            for i in range(self.N):
                self.train_part += RZ(f'rz_{p}{i}').on(i)

            # generate U_B
            for i in range(self.N):
                self.train_part += RX(f'beta_{p}{i}').on(i)
        self.circ = self.init_part + self.train_part
        self.hs = np.zeros(self.N)
        # Generate a string like 'Z0 Z1 Z2 ... Z{N-1}', with a single space between each element
        # self.z_ops = ' '.join([f'Z{i}' for i in range(self.N)])
    
        if self.add_noise:
            from mindquantum.core.circuit import MixerAdder, DepolarizingChannelAdder, QubitNumberConstrain
            # we add depolarizing noise to the circuit
            p1, p2 = self.noise_param
            noise_adder1 = MixerAdder([QubitNumberConstrain(1),DepolarizingChannelAdder(p1, 1),], add_after=True)
            noise_adder2 = MixerAdder([QubitNumberConstrain(2),DepolarizingChannelAdder(p2, 2),], add_after=True)

            self.circ = noise_adder2(noise_adder1(self.circ))

    def _build_ham(self) -> np.ndarray:
        ham_matrix = get_ham_matrix(self.config['BIG_adjacent_weight_matrix'], self.config['BIG_adjacent_matrix'], self.config['P'])
        return ham_matrix

    def _update_hs(self):
        sim_hs = Simulator('mqvector', self.circ.n_qubits)
        # pr = dict(zip(self.circ.params_name, self.thetas))
        pr = dict(zip(self.circ.params_name, np.concatenate((2*np.arctan(1/(np.sqrt(1 + self.hs**2) - self.hs)), self.thetas))))
        # print(self.thetas)
        for i in range(self.N):
            sim_hs.reset()
            sim_hs.apply_circuit(self.circ, pr=pr)
            z_exp = sim_hs.get_expectation(Hamiltonian(QubitOperator(f'Z{i}')))
            self.hs[i] = self.hs[i] + self.config['abf_lr'] * (z_exp.real - self.hs[i])

    def evolution(self, print_ratio: bool = False) -> np.ndarray:
        lr = self.config['lr']
        epoch = self.config['epoch']
 
        sim = Simulator('mqvector', self.circ.n_qubits)                     # 创建模拟器，backend使用‘mqvector’，能模拟5个比特（'circ'线路中包含的比特数）
        
        ham_matrix = self._build_ham()
        ham = csr_matrix(ham_matrix)

        min_eigenvalue, min_eigenvector = sla.eigsh(ham, k=1, which='SA')
        max_eigenvalue, max_eigenvector = sla.eigsh(ham, k=1, which='LA')

        grad_ops = sim.get_expectation_with_grad(Hamiltonian(ham), self.train_part)            # 获取计算变分量子线路的期望值和梯度的算子
        net = MQAnsatzOnlyLayer(grad_ops)                              # 生成待训练的神经网络
        opti = nn.Adam(net.trainable_params(), learning_rate=lr)     # 设置针对网络中所有可训练参数、学习率为lr的Adam优化器
        train_net = nn.TrainOneStepCell(net, opti)

        self.training_data = []
        from tqdm import tqdm  # import tqdm for progress bar

        for e in tqdm(range(epoch), desc="Training Progress"):
            train_net()
            thetas = net.weight.asnumpy()
            # thetas = np.concatenate((2*np.arctan(1/(np.sqrt(1 + self.hs**2) - self.hs)), thetas[self.N:]))
            # print(thetas)
            self.thetas = thetas
            f, _ = grad_ops(thetas)
            ratio = (max_eigenvalue-f.real)/(max_eigenvalue-min_eigenvalue)
            self.training_data.append(ratio[0][0])
            if print_ratio and e % 10 == 0:
                print(f'epoch: {e}, approximate ratio: {ratio[0][0]}')
            self._update_hs()
        
        print(f'final approximate ratio: {ratio[0][0]}')
        print('ABF-QAOA parameter optimization finished')
        
    
    def get_successful_ratio(self, shots: int, show_result: bool = False) -> float:
        successful_ratio = 0.0
        
        circ = self.circ
        circ.measure_all()
        sim = Simulator('mqvector', circ.n_qubits)
        
        pr = dict(zip(circ.params_name, np.concatenate((self.hs, self.thetas))))
        # pr = dict(zip(circ.params_name, self.thetas))
        result = sim.sampling(circ, pr=pr, shots=shots).data # 比特串采用小端序表示(返回字典)
        if show_result:
            print(result)
        if self.ideal_state in result:
            successful_ratio = result[self.ideal_state]/shots
        else:
            print(f'ideal state {self.ideal_state} not found in result')
        
        return successful_ratio
    
    def get_training_data(self) -> np.ndarray:
        return np.array(self.training_data)